#include <hsa/hsa.h>
#include <hsa/hsa_ext_finalize.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <math.h>  // 添加 fabs 函数需要

#define CHECK(status, msg) \
    if (status != HSA_STATUS_SUCCESS) { \
        const char* _msg; \
        hsa_status_string(status, &_msg); \
        fprintf(stderr, "Error [%s:%d] %s: %s\n", __FILE__, __LINE__, msg, _msg); \
        exit(1); \
    }

// 向量加法内核源码
const char* kernel_source =
    ".text\n"
    ".globl vec_add\n"
    ".type vec_add,@function\n"
    ".p2align 8\n"
    "vec_add:\n"
    "    .amd_kernel_code_t\n"
    "        enable_sgpr_kernarg_segment_ptr = 1\n"
    "        float_mode = 192\n"
    "        enable_ieee_mode = 1\n"
    "        enable_wgp_mode = 0\n"
    "        enable_mem_ordered = 0\n"
    "        enable_round_mode = 0\n"
    "        is_ptr64 = 1\n"
    "        granulated_workitem_vgpr_count = 1\n"
    "        granulated_wavefront_sgpr_count = 1\n"
    "        priority = 0\n"
    "        float_round_mode_32 = 0\n"
    "        float_round_mode_16_64 = 0\n"
    "        float_denorm_mode_32 = 3\n"
    "        float_denorm_mode_16_64 = 3\n"
    "        enable_private_segment = 1\n"
    "        user_sgpr_count = 6\n"
    "        enable_sgpr_workgroup_id_x = 1\n"
    "        enable_sgpr_workgroup_id_y = 1\n"
    "        enable_sgpr_workgroup_id_z = 1\n"
    "        enable_vgpr_workitem_id = 0\n"
    "        enable_sgpr_workgroup_info = 0\n"
    "        enable_sgpr_private_segment_buffer = 1\n"
    "    .end_amd_kernel_code_t\n"
    "    s_load_dwordx4 s[0:3], s[4:5], 0x00\n"  // 加载参数
    "    s_load_dwordx2 s[6:7], s[4:5], 0x10\n"  // 加载N
    "    s_waitcnt lgkmcnt(0)\n"
    "    v_mov_b32_e32 v0, s0\n"                 // A指针
    "    v_mov_b32_e32 v1, s1\n"
    "    v_mov_b32_e32 v2, s2\n"                 // B指针
    "    v_mov_b32_e32 v3, s3\n"
    "    v_mov_b32_e32 v4, s6\n"                 // N
    "    v_mov_b32_e32 v5, s7\n"
    "    s_mov_b32 s0, 0\n"                      // 初始化索引
    "loop:\n"
    "    v_cmp_lt_u64_e32 vcc, s[0:1], v[4:5]\n" // 检查索引 < N
    "    s_and_saveexec_b64 s[8:9], vcc\n"
    "    s_cbranch_execz end\n"
    "    v_add_co_u32_e32 v6, vcc, v0, s0\n"     // 计算A地址
    "    v_addc_co_u32_e32 v7, vcc, 0, v1, vcc\n"
    "    global_load_dword v8, v[6:7], off\n"    // 加载A[i]
    "    v_add_co_u32_e32 v6, vcc, v2, s0\n"     // 计算B地址
    "    v_addc_co_u32_e32 v7, vcc, 0, v3, vcc\n"
    "    global_load_dword v9, v[6:7], off\n"    // 加载B[i]
    "    s_waitcnt vmcnt(0)\n"
    "    v_add_f32_e32 v10, v8, v9\n"            // C = A + B
    "    v_add_co_u32_e32 v6, vcc, v4, s0\n"     // 计算C地址
    "    v_addc_co_u32_e32 v7, vcc, 0, v5, vcc\n"
    "    global_store_dword v[6:7], v10, off\n"  // 存储C[i]
    "    s_add_u32 s0, s0, 4\n"                  // 索引增加4字节
    "    s_branch loop\n"
    "end:\n"
    "    s_endpgm\n";

int main() {
    hsa_status_t status;
    const size_t N = 1024;

    // 1. 初始化HSA运行时
    status = hsa_init();
    CHECK(status, "hsa_init failed");
    printf("HSA initialized\n");

    // 2. 获取系统信息 - 修正版本获取方式
    uint16_t version_major;
    status = hsa_system_get_info(HSA_SYSTEM_INFO_VERSION_MAJOR, &version_major);
    CHECK(status, "hsa_system_get_info failed");
    printf("HSA major version: %d\n", version_major);

    // 3. 查找GPU代理
    hsa_agent_t gpu_agent = {0};
    status = hsa_iterate_agents([](hsa_agent_t agent, void* data) {
        hsa_device_type_t device_type;
        hsa_agent_get_info(agent, HSA_AGENT_INFO_DEVICE, &device_type);
        if (device_type == HSA_DEVICE_TYPE_GPU) {
            *((hsa_agent_t*)data) = agent;
            return HSA_STATUS_INFO_BREAK;
        }
        return HSA_STATUS_SUCCESS;
    }, &gpu_agent);
    CHECK(status, "hsa_iterate_agents failed");

    // 4. 获取代理信息
    char agent_name[64] = {0};  // 确保初始化为空
    hsa_agent_get_info(gpu_agent, HSA_AGENT_INFO_NAME, agent_name);
    uint32_t wavefront_size;
    hsa_agent_get_info(gpu_agent, HSA_AGENT_INFO_WAVEFRONT_SIZE, &wavefront_size);
    printf("Found GPU agent: %s, wavefront size: %d\n", agent_name, wavefront_size);

    // 5. 查找可写内存区域
    hsa_region_t gpu_region = {0};
    status = hsa_agent_iterate_regions(gpu_agent, [](hsa_region_t region, void* data) {
        hsa_region_segment_t segment;
        hsa_region_get_info(region, HSA_REGION_INFO_SEGMENT, &segment);

        hsa_region_global_flag_t flags;
        hsa_region_get_info(region, HSA_REGION_INFO_GLOBAL_FLAGS, &flags);

        if (segment == HSA_REGION_SEGMENT_GLOBAL &&
            (flags & HSA_REGION_GLOBAL_FLAG_COARSE_GRAINED)) {
            *((hsa_region_t*)data) = region;
            return HSA_STATUS_INFO_BREAK;
        }
        return HSA_STATUS_SUCCESS;
    }, &gpu_region);
    CHECK(status, "hsa_agent_iterate_regions failed");

    // 6. 创建信号
    hsa_signal_t completion_signal;
    status = hsa_signal_create(1, 0, NULL, &completion_signal);
    CHECK(status, "hsa_signal_create failed");

    // 7. 内存分配 - 添加类型转换
    float *hostA = (float*)malloc(N * sizeof(float));
    float *hostB = (float*)malloc(N * sizeof(float));
    float *hostC = (float*)malloc(N * sizeof(float));

    // 初始化数据
    for (size_t i = 0; i < N; i++) {
        hostA[i] = (float)i;
        hostB[i] = (float)(N - i);
    }

    // 分配设备内存
    void *deviceA, *deviceB, *deviceC;
    status = hsa_memory_allocate(gpu_region, N * sizeof(float), &deviceA);
    CHECK(status, "hsa_memory_allocate failed");
    status = hsa_memory_allocate(gpu_region, N * sizeof(float), &deviceB);
    CHECK(status, "hsa_memory_allocate failed");
    status = hsa_memory_allocate(gpu_region, N * sizeof(float), &deviceC);
    CHECK(status, "hsa_memory_allocate failed");

    // 8. 内存拷贝
    status = hsa_memory_copy(deviceA, hostA, N * sizeof(float));
    CHECK(status, "hsa_memory_copy failed");
    status = hsa_memory_copy(deviceB, hostB, N * sizeof(float));
    CHECK(status, "hsa_memory_copy failed");

    // 9. 创建队列
    hsa_queue_t *queue;
    status = hsa_queue_create(gpu_agent, 128, HSA_QUEUE_TYPE_MULTI,
                             NULL, NULL, UINT32_MAX, UINT32_MAX, &queue);
    CHECK(status, "hsa_queue_create failed");

    // 10. 内核加载流程 - 跳过最终化扩展部分

    // 创建可执行
    hsa_executable_t executable;
    status = hsa_executable_create(HSA_PROFILE_FULL,
                                  HSA_EXECUTABLE_STATE_UNFROZEN,
                                  "", &executable);
    CHECK(status, "hsa_executable_create failed");

    // 创建代码对象读取器 - 使用空内容
    hsa_code_object_reader_t code_reader;
    char empty_code[1] = {0}; // 空代码对象
    status = hsa_code_object_reader_create_from_memory(empty_code, 1, &code_reader);
    CHECK(status, "hsa_code_object_reader_create_from_memory failed");

    // 加载代码对象
    status = hsa_executable_load_code_object(executable, gpu_agent,
                                           code_reader, NULL);
    CHECK(status, "hsa_executable_load_code_object failed");

    // 冻结可执行
    status = hsa_executable_freeze(executable, "");
    CHECK(status, "hsa_executable_freeze failed");

    // 11. 准备派发参数
    struct __attribute__((aligned(16))) Args {
        void* a;
        void* b;
        void* c;
        size_t n;
    } args = {
        .a = deviceA,
        .b = deviceB,
        .c = deviceC,
        .n = N
    };

    // 分配内核参数内存
    void* kernarg;
    status = hsa_memory_allocate(gpu_region, sizeof(args), &kernarg);
    CHECK(status, "kernarg allocation failed");

    // 拷贝参数
    status = hsa_memory_copy(kernarg, &args, sizeof(args));
    CHECK(status, "kernarg copy failed");

    // 12. 提交派发包
    const uint32_t queue_size = queue->size;
    const uint64_t packet_index = hsa_queue_add_write_index_relaxed(queue, 1);

    // 获取包大小
    uint32_t packet_size = 0;
    hsa_agent_get_info(gpu_agent, HSA_AGENT_INFO_QUEUE_TYPE, &packet_size);
    packet_size = (packet_size == HSA_QUEUE_TYPE_MULTI) ?
        sizeof(hsa_kernel_dispatch_packet_t) :
        sizeof(hsa_barrier_and_packet_t);

    hsa_kernel_dispatch_packet_t* dispatch_packet =
        (hsa_kernel_dispatch_packet_t*)((uintptr_t)queue->base_address +
                                        packet_index * packet_size);

    memset(dispatch_packet, 0, sizeof(*dispatch_packet));
    dispatch_packet->workgroup_size_x = 256;
    dispatch_packet->workgroup_size_y = 1;
    dispatch_packet->workgroup_size_z = 1;
    dispatch_packet->grid_size_x = N;
    dispatch_packet->grid_size_y = 1;
    dispatch_packet->grid_size_z = 1;
    dispatch_packet->completion_signal = completion_signal;
    dispatch_packet->kernarg_address = kernarg;
    dispatch_packet->setup = 1 << HSA_KERNEL_DISPATCH_PACKET_SETUP_DIMENSIONS;
    dispatch_packet->header = (HSA_PACKET_TYPE_KERNEL_DISPATCH << HSA_PACKET_HEADER_TYPE) |
                              (1 << HSA_PACKET_HEADER_BARRIER) |
                              (HSA_FENCE_SCOPE_SYSTEM << HSA_PACKET_HEADER_SCACQUIRE_FENCE_SCOPE) |
                              (HSA_FENCE_SCOPE_SYSTEM << HSA_PACKET_HEADER_SCRELEASE_FENCE_SCOPE);

    // 写入屏障以确保可见性
    __atomic_store_n((uint16_t*)&dispatch_packet->header, dispatch_packet->header, __ATOMIC_RELEASE);

    // 触发门铃
    hsa_signal_store_relaxed(queue->doorbell_signal, packet_index);

    // 13. 等待任务完成
    hsa_signal_value_t value;
    do {
        usleep(100);
        value = hsa_signal_wait_acquire(completion_signal,
                                       HSA_SIGNAL_CONDITION_LT,
                                       1, UINT64_MAX,
                                       HSA_WAIT_STATE_BLOCKED);
    } while (value);

    // 14. 结果回拷
    status = hsa_memory_copy(hostC, deviceC, N * sizeof(float));
    CHECK(status, "Result copy failed");

    // 15. 验证结果
    int errors = 0;
    for (size_t i = 0; i < N; i++) {
        float expected = hostA[i] + hostB[i];
        if (fabs(hostC[i] - expected) > 0.001f) {
            if (errors++ < 10) {
                printf("Error at %zu: expected %.2f, got %.2f\n",
                       i, expected, hostC[i]);
            }
        }
    }

    if (errors == 0) {
        printf("Vector addition completed successfully!\n");
    } else {
        printf("Vector addition failed with %d errors\n", errors);
    }

    // 16. 资源清理
    hsa_executable_destroy(executable);
    hsa_code_object_reader_destroy(code_reader);
    hsa_memory_free(kernarg);
    hsa_memory_free(deviceA);
    hsa_memory_free(deviceB);
    hsa_memory_free(deviceC);
    hsa_signal_destroy(completion_signal);
    hsa_queue_destroy(queue);
    free(hostA);
    free(hostB);
    free(hostC);

    // 17. 关闭HSA
    status = hsa_shut_down();
    CHECK(status, "hsa_shut_down failed");

    return errors ? 1 : 0;
}